import torch
import torch.nn as nn

import numpy as np
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
import time


checkpoint = torch.load("/media/data1/wjh/MetaR-CNN/meta_class_agnostic/third/pascal_voc_0712_metarcnn_1_30_51.pth")
model = checkpoint['model']

for key, value in model.items():
    if key in ['RCNN_cls_score.weight']:
        print('key is : %s | value shape is : %s' % (key, value.shape))
        value1 = value.cpu().numpy()
        for cls in range(21):
            print('%.0f norm is : %.3f' % (cls, np.linalg.norm(value1[cls, :])))
            #np.savetxt('weight/%s.csv' % key, value, fmt='%.3f', delimiter=',')
            

'''
print(value1.shape)
data = value1

t1 = time.time()
tsne = TSNE(n_components = 2).fit_transform(data)
print('tsne time used: %.2f' % float(time.time() - t1))


color = ['#FFD700',
'#808080',
'#00FFFF',
'#7FFFD4',
'#008000',
'#90EE90',
'#D3D3D3',
'#000000',
'#FFEBCD',
'#0000FF',
'#8A2BE2',
'#A52A2A',
'#DEB887',
'#5F9EA0',
'#7FFF00',
'#D2691E',
'#FF7F50',
'#6495ED',
'#FFA500',
'#DC143C',
'#FFD700']

for cls in range(21):

    x = tsne[cls, 0]
    y = tsne[cls, 1]

    plt.scatter(x, y, s=1, c = color[cls])

plt.legend(['bicycle', 'bird', 'boat', 'bus', 'car', 'cat', 'chair', 'diningtable', 'dog', 'motorbike',
                         'person', 'pottedplant', 'sheep', 'train', 'tvmonitor', 'aeroplane', 'bottle', 'cow',
                         'horse', 'sofa'], bbox_to_anchor=(0.95, 1))
                         
plt.savefig('tsne_weight.jpg', dpi=1000)
'''